#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: pycharm
@file: cluster_bboxes.py
@time: 2020/9/24 15:30
@desc:
'''
import numpy as np

def pt_in_rect(rect1,rect2):
    '''
    判断rect1的中心点是否落在rect2内部
    :param rect1: xmin,ymin,xmax,ymax
    :param rect2:
    :return:
    '''
    xc = (rect1[0] + rect1[2]) / 2
    yc = (rect1[1] + rect1[3]) / 2
    xmin,ymin,xmax,ymax = rect2
    return xc >= xmin and xc <= xmax and yc >= ymin and yc <= ymax

def judge_same_col(rect1,rect2):
    '''
    判断两个矩形是否属于同一列
    :param rect1: xmin,ymin,xmax,ymax
    :param rect2:
    :return:
    '''
    x1, _, x2, *_ = rect1
    xc = (x1 + x2) / 2
    x1, _, x2, *_ = rect2
    return xc >= x1 and xc <= x2

def judge_same_row(rect1,rect2):
    '''
    判断两个矩形是否属于同一行
    :param rect1: xmin,ymin,xmax,ymax
    :param rect2:
    :return:
    '''
    _, y1, _, y2, *_ = rect1
    yc = (y1 + y2) / 2
    _, y1, _, y2, *_ = rect2
    return yc >= y1 and yc <= y2

def cluster_bboxes(bboxes,direction=0):
    '''
    基于列聚类
    :param bboxes:
    :param direction:0为按列聚类，1为按行聚类
    :return:
    '''
    bboxes_clustered_x = []  # 以列的形式存放，每个列表代表一列
    if direction == 0:
        sorted_bboxes_x = sorted(bboxes, key=lambda bbox: bbox[0] + bbox[2], reverse=True)
        judge_func = judge_same_col
    elif direction == 1:
        sorted_bboxes_x = sorted(bboxes, key=lambda bbox: bbox[1] + bbox[3], reverse=True)
        judge_func = judge_same_row
    #按x排序按列聚类
    while True:
        bboxes_removed = []
        cur_box = sorted_bboxes_x.pop()
        bboxes_removed.append(cur_box)
        for i in range(len(sorted_bboxes_x)-1,-1,-1):
            next_box = sorted_bboxes_x[i]
            if judge_func(next_box,cur_box):
                cur_box = next_box
                bboxes_removed.append(cur_box)
                del sorted_bboxes_x[i]
            else:
                break
        bboxes_clustered_x.append(bboxes_removed)
        if len(sorted_bboxes_x)==0:
            break
    for i in range(len(bboxes_clustered_x)):
        if direction == 0:
            bboxes_col = sorted(bboxes_clustered_x[i],key=lambda bbox:bbox[1]+bbox[3])
        elif direction == 1:
            bboxes_col = sorted(bboxes_clustered_x[i], key=lambda bbox: bbox[0] + bbox[2])
        bboxes_clustered_x[i] = bboxes_col
    return bboxes_clustered_x


def my_tolist(targets):
    if not isinstance(targets,np.ndarray):
        results = []
        for target in targets:
            results.append(my_tolist(target))
    else:
        results = targets.tolist()
    return results


def cluster(bboxes):
    # 计算所有矩形框的平均宽高
    avr_width = 0
    avr_height = 0
    for xmin, ymin, xmax, ymax, *_ in bboxes:
        avr_width += xmax - xmin
        avr_height += ymax - ymin
    avr_width /= len(bboxes)
    avr_height /= len(bboxes)

    # 所有框按照行、列进行聚类
    bboxes_clustered_x = my_tolist(cluster_bboxes(bboxes, direction=0))
    bboxes_clustered_y = my_tolist(cluster_bboxes(bboxes, direction=1))

    cols = len(bboxes_clustered_x)
    rows = len(bboxes_clustered_y)
    print(rows, cols)

    # 所有框进行序列化输出
    serialized_output = []
    for r in range(rows):
        tmp = []
        for c in range(cols):
            tmp.append([])
        serialized_output.append(tmp)
    for bbox in bboxes:
        bbox = bbox.tolist()
        loc_r = -1
        loc_c = -1
        for r in range(len(bboxes_clustered_y)):
            if bbox in bboxes_clustered_y[r]:
                loc_r = r
                break
        for c in range(len(bboxes_clustered_x)):
            if bbox in bboxes_clustered_x[c]:
                loc_c = c
                break
        if loc_r != -1 and loc_c != -1:
            if (len(serialized_output[loc_r][loc_c]) == 0) or (serialized_output[loc_r][loc_c][-2] <= bbox[-2]):
                serialized_output[loc_r][loc_c] = bbox
        else:
            raise NotImplementedError

    # 对未检测到的位置进行补齐
    for r in range(rows):
        for c in range(cols):
            if len(serialized_output[r][c]) == 0:
                cur_col = np.array(bboxes_clustered_x[c])
                avr_x = np.mean(cur_col[:, 0])
                cur_row = np.array(bboxes_clustered_y[r])
                avr_y = np.mean(cur_row[:, 1])
                serialized_output[r][c] = [avr_x, avr_y, avr_x + avr_width, avr_y + avr_height, 0, 3]

    for output in serialized_output:
        print(output)






bboxes = np.random.randint(1,20,(10,6))
bboxes[:,2] += bboxes[:,0]
bboxes[:,3] += bboxes[:,1]
print(bboxes)
cluster(bboxes)